/*
 * Copyright 2023 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.firebase.ai.java

import androidx.concurrent.futures.SuspendToFutureAdapter
import com.google.common.util.concurrent.ListenableFuture
import com.google.firebase.ai.GenerativeModel
import com.google.firebase.ai.java.ChatFutures.Companion.from
import com.google.firebase.ai.type.Content
import com.google.firebase.ai.type.CountTokensResponse
import com.google.firebase.ai.type.FirebaseAIException
import com.google.firebase.ai.type.GenerateContentResponse
import kotlinx.coroutines.reactive.asPublisher
import org.reactivestreams.Publisher

/**
 * Wrapper class providing Java compatible methods for [GenerativeModel].
 *
 * @see [GenerativeModel]
 */
public abstract class GenerativeModelFutures internal constructor() {

  /**
   * Generates new content from the input [Content] given to the model as a prompt.
   *
   * @param prompt The input(s) given to the model as a prompt.
   * @return The content generated by the model.
   * @throws [FirebaseAIException] if the request failed.
   */
  public abstract fun generateContent(
    prompt: Content,
    vararg prompts: Content
  ): ListenableFuture<GenerateContentResponse>

  /**
   * Generates new content as a stream from the input [Content] given to the model as a prompt.
   *
   * @param prompt The input(s) given to the model as a prompt.
   * @return A [Publisher] which will emit responses as they are returned by the model.
   * @throws [FirebaseAIException] if the request failed.
   */
  public abstract fun generateContentStream(
    prompt: Content,
    vararg prompts: Content
  ): Publisher<GenerateContentResponse>

  /**
   * Counts the number of tokens in a prompt using the model's tokenizer.
   *
   * @param prompt The input(s) given to the model as a prompt.
   * @return The [CountTokensResponse] of running the model's tokenizer on the input.
   * @throws [FirebaseAIException] if the request failed.
   */
  public abstract fun countTokens(
    prompt: Content,
    vararg prompts: Content
  ): ListenableFuture<CountTokensResponse>

  /**
   * Creates a [ChatFutures] instance which internally tracks the ongoing conversation with the
   * model.
   */
  public abstract fun startChat(): ChatFutures

  /**
   * Creates a [ChatFutures] instance, initialized using the optionally provided [history].
   *
   * @param history A list of previous interactions with the model to use as a starting point
   */
  public abstract fun startChat(history: List<Content>): ChatFutures

  /** Returns the [GenerativeModel] object wrapped by this object. */
  public abstract fun getGenerativeModel(): GenerativeModel

  private class FuturesImpl(private val model: GenerativeModel) : GenerativeModelFutures() {
    override fun generateContent(
      prompt: Content,
      vararg prompts: Content
    ): ListenableFuture<GenerateContentResponse> =
      SuspendToFutureAdapter.launchFuture { model.generateContent(prompt, *prompts) }

    override fun generateContentStream(
      prompt: Content,
      vararg prompts: Content
    ): Publisher<GenerateContentResponse> =
      model.generateContentStream(prompt, *prompts).asPublisher()

    override fun countTokens(
      prompt: Content,
      vararg prompts: Content
    ): ListenableFuture<CountTokensResponse> =
      SuspendToFutureAdapter.launchFuture { model.countTokens(prompt, *prompts) }

    override fun startChat(): ChatFutures = startChat(emptyList())

    override fun startChat(history: List<Content>): ChatFutures = from(model.startChat(history))

    override fun getGenerativeModel(): GenerativeModel = model
  }

  public companion object {

    /** @return a [GenerativeModelFutures] created around the provided [GenerativeModel] */
    @JvmStatic public fun from(model: GenerativeModel): GenerativeModelFutures = FuturesImpl(model)
  }
}
